import akshare as ak
import pandas as pd
import numpy as np
import os
import time
from datetime import datetime, date
from dateutil.relativedelta import relativedelta
import traceback

from date_utils import AShareTradingDayChecker, DeliveryDateCalculator
from components import DatePicker, Logger

class data_fetcher:
    """数据获取器，封装所有数据获取逻辑"""
    def __init__(self, cache_dir="stock_cache", trading_day_checker=None):
        """
        初始化数据获取器
        :param cache_dir: 缓存目录路径
        :param trading_day_checker: 交易日检查器实例
        """
        self.cache_dir = cache_dir
        self.trading_day_checker = trading_day_checker

        # 第一步：初始化日志系统
        self.logger = Logger()  # 先初始化logger

        self.turnover_cache_path = os.path.join(self.cache_dir, "turnover_cache.csv") 
        
        # 确保缓存目录存在
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)
    
    def get_market_index_data(self, index_code, start_date, end_date):
        """获取指数数据（日期类型安全版）"""       
        try:
            # 尝试从缓存加载
            cache_file = os.path.join(self.cache_dir, f"index_{index_code}.csv")
            if os.path.exists(cache_file):
                df = pd.read_csv(cache_file, parse_dates=['date'])
                df['date'] = pd.to_datetime(df['date'])  # 强制类型转换
            else:
                # 从API获取新数据
                df = ak.stock_zh_index_daily(symbol=index_code)
                df['date'] = pd.to_datetime(df['date'])  # 确保日期列类型
            
            # 日期范围过滤
            mask = (df['date'] >= start_date) & (df['date'] <= end_date)
            filtered_df = df.loc[mask].copy()
            
            if filtered_df.empty:
                raise ValueError(f"无有效数据：{index_code} ({start_date.date()}~{end_date.date()})")
                
            # 计算涨跌幅
            filtered_df['pct_change'] = filtered_df['close'].pct_change() * 100
            return filtered_df

        except Exception as e:
            self.logger.log(f"指数数据加载失败：{traceback.format_exc()}")
            return pd.DataFrame(columns=['date', 'close', 'pct_change'])
    
    def get_north_flow_data(self, start_date, end_date):
        """获取北向资金净流入数据（兼容多种akshare版本）"""
        try:
            # 尝试最新接口
            try:
                df = ak.stock_hsgt_north_net_flow_in_em()
                if df is None or df.empty:
                    raise Exception("获取北向资金数据失败")
                    
                df['date'] = pd.to_datetime(df['date'])
                df = df[(df['date'] >= start_date) & (df['date'] <= end_date)]
                df = df.rename(columns={'net_amount': 'north_net_inflow'})
                # 转换为亿元单位
                df['north_net_inflow'] = df['north_net_inflow'] / 100000000
                
                self.logger.log(f"使用新接口获取北向资金数据成功，共{len(df)}条记录")
                return df[['date', 'north_net_inflow']]
            except Exception as e1:
                self.logger.log(f"尝试最新北向资金接口失败: {e1}")
                
                # 尝试旧接口
                try:
                    df = ak.stock_individual_fund_flow_rank()
                    if df is None or df.empty:
                        raise Exception("获取北向资金数据失败")
                        
                    df['date'] = pd.to_datetime(df['日期'])
                    df = df[(df['date'] >= start_date) & (df['date'] <= end_date)]
                    df = df.rename(columns={'主力净流入-净额': 'north_net_inflow'})
                    
                    self.logger.log(f"使用旧接口1获取北向资金数据成功，共{len(df)}条记录")
                    return df[['date', 'north_net_inflow']]
                except Exception as e2:
                    self.logger.log(f"尝试旧北向资金接口失败: {e2}")
                    
                    # 尝试备选接口
                    try:
                        df = ak.stock_market_fund_flow()
                        if df is None or df.empty:
                            raise Exception("获取北向资金数据失败")
                            
                        df['date'] = pd.to_datetime(df['日期'])
                        df = df[(df['date'] >= start_date) & (df['date'] <= end_date)]
                        df = df.rename(columns={'主力净流入-净额': 'north_net_inflow'})
                        
                        self.logger.log(f"使用旧接口2获取北向资金数据成功，共{len(df)}条记录")
                        return df[['date', 'north_net_inflow']]
                    except Exception as e3:
                        self.logger.log(f"尝试备选北向资金接口失败: {e3}")
                        raise Exception("无法获取北向资金数据，请确保akshare版本是最新的")
        except Exception as e:
            self.logger.log(f"获取北向资金数据时出错: {e}")
            return None

    def get_main_money_flow_data(self, start_date, end_date):
        """获取主力资金数据（带本地缓存）"""
        filename = os.path.join(self.cache_dir, "main_money_flow.csv")
        local_data = pd.DataFrame()

        # 尝试读取本地缓存
        if os.path.exists(filename):
            try:
                local_data = pd.read_csv(filename, parse_dates=['date'])
                local_data = local_data[(local_data['date'] >= start_date) & 
                                    (local_data['date'] <= end_date)]
                self.logger.log(f"从本地缓存加载主力资金数据（{len(local_data)} 条）")
            except Exception as e:
                self.logger.log(f"读取主力资金缓存失败: {e}")

        # 检查本地数据是否满足需求
        if not local_data.empty:
            cache_start = local_data['date'].min()
            cache_end = local_data['date'].max()
            if cache_start <= start_date and cache_end >= end_date:
                return local_data

        # 从API获取数据（兼容多个接口）
        api_data = pd.DataFrame()
        try:
            # 尝试stock_market_fund_flow接口
            try:
                df = ak.stock_market_fund_flow()
                if df is not None and not df.empty:
                    df['date'] = pd.to_datetime(df['日期'])
                    df = df.rename(columns={'主力净流入-净额': 'main_net_inflow'})
                    # 转换单位（假设原始数据为元）
                    df['main_net_inflow'] = df['main_net_inflow'] / 100000000
                    api_data = df[['date', 'main_net_inflow']]
            except Exception as e:
                self.logger.log(f"接口1获取失败: {e}")

            # 如果第一个接口失败，尝试stock_individual_fund_flow
            if api_data.empty:
                try:
                    df = ak.stock_individual_fund_flow(stock="all")
                    if df is not None and not df.empty:
                        # 按日期聚合
                        df['date'] = pd.to_datetime(df['日期'])
                        grouped = df.groupby('date')['主力净流入-净额'].sum().reset_index()
                        grouped = grouped.rename(columns={'主力净流入-净额': 'main_net_inflow'})
                        # 转换单位
                        grouped['main_net_inflow'] = grouped['main_net_inflow'] / 100000000
                        api_data = grouped[['date', 'main_net_inflow']]
                except Exception as e:
                    self.logger.log(f"接口2获取失败: {e}")

            # 过滤API数据日期范围
            if not api_data.empty:
                api_data = api_data[(api_data['date'] >= start_date) & 
                                (api_data['date'] <= end_date)]
                self.logger.log(f"从API获取主力资金数据 {len(api_data)} 条")

                # 合并本地缓存和API数据
                combined = pd.concat([local_data, api_data])
                combined = combined.drop_duplicates('date').sort_values('date')

                # 保存更新后的缓存
                try:
                    # 读取完整缓存数据
                    full_cache = pd.DataFrame()
                    if os.path.exists(filename):
                        full_cache = pd.read_csv(filename, parse_dates=['date'])
                    
                    # 合并新旧数据
                    updated_cache = pd.concat([full_cache, combined])
                    updated_cache = updated_cache.drop_duplicates('date')
                    updated_cache.to_csv(filename, index=False)
                    self.logger.log(f"主力资金缓存已更新，总数据量 {len(updated_cache)} 条")
                except Exception as e:
                    self.logger.log(f"保存主力资金缓存失败: {e}")

                return combined

            return local_data  # 返回本地缓存数据（即使不完整）
        except Exception as e:
            self.logger.log(f"获取主力资金数据失败: {str(e)}")
            return local_data  # 返回本地缓存数据（如果有）

    def get_market_turnover(self, start_date, end_date):
        """获取成交额数据（修复日期类型问题）"""
        try:
            # === 关键修复：保持日期为 Pandas Timestamp 类型 ===
            start_date = pd.to_datetime(start_date)
            end_date = pd.to_datetime(end_date)
            
            # 尝试从缓存加载
            cached_data = self._load_turnover_cache()
            
            # 检查缓存是否满足需求
            if not cached_data.empty:
                cache_start = cached_data['date'].min()
                cache_end = cached_data['date'].max()
                
                if cache_start <= start_date and cache_end >= end_date:
                    mask = (cached_data['date'] >= start_date) & (cached_data['date'] <= end_date)
                    return cached_data.loc[mask].copy()
            
            # 需要获取新数据
            new_data = self._fetch_full_turnover_data(start_date, end_date)
            
            # 合并新旧数据（自动处理日期类型）
            combined = pd.concat([cached_data, new_data])
            combined = combined.drop_duplicates('date').sort_values('date')
            
            # 保存更新后的缓存
            self._save_turnover_cache(combined)
            
            # 过滤请求范围
            mask = (combined['date'] >= start_date) & (combined['date'] <= end_date)
            return combined.loc[mask].copy()
            
        except Exception as e:
            self.logger.log(f"获取成交额失败: {str(e)}")
            return pd.DataFrame(columns=['date', 'total_turnover'])
        

    def _load_turnover_cache(self):
        """加载成交额缓存（修复日期解析）"""
        try:
            if os.path.exists(self.turnover_cache_path):
                df = pd.read_csv(
                    self.turnover_cache_path,
                    parse_dates=['date'],  # 自动解析日期列为 datetime64[ns]
                    date_parser=lambda x: pd.to_datetime(x, format='%Y-%m-%d')
                )
                # 验证数据列
                required_cols = ['date', 'total_turnover']
                if all(col in df.columns for col in required_cols):
                    return df[required_cols].sort_values('date')
            return pd.DataFrame(columns=['date', 'total_turnover'])
        except Exception as e:
            self.logger.log(f"加载成交额缓存失败: {str(e)}")
            return pd.DataFrame(columns=['date', 'total_turnover'])

    def _save_turnover_cache(self, data):
        """保存数据到缓存文件"""
        try:
            # 保留最近3年数据
            cutoff_date = pd.Timestamp.now() - pd.DateOffset(years=3)
            filtered = data[data['date'] > cutoff_date]
            
            filtered.to_csv(self.turnover_cache_path, index=False)
            self.logger.log(f"成交额缓存已更新，保存{len(filtered)}条记录")
        except Exception as e:
            self.logger.log(f"保存成交额缓存失败: {str(e)}")


    def _fetch_full_turnover_data(self, start_date: date, end_date: date) -> pd.DataFrame:
        """获取完整成交额数据（修复变量初始化问题）"""
        try:
            # 初始化数据容器
            sse_data = []
            szse_data = []
            dates_str = []

            # 获取交易日列表
            trade_dates = self.trading_day_checker.get_trading_days(start_date, end_date)
            if not trade_dates:
                self.logger.log(f"无交易日数据 {start_date} 至 {end_date}")
                return pd.DataFrame(columns=['date', 'sse_turnover', 'szse_turnover', 'total_turnover'])

            # 准备日期列表
            dates_str = [d.strftime("%Y%m%d") for d in trade_dates]

            # 获取交易所数据
            for date_str in dates_str:
                try:
                    sse = self.fetch_sse_turnover(date_str)
                    szse = self.fetch_szse_turnover(date_str)
                    sse_data.append(sse)
                    szse_data.append(szse)
                except Exception as e:
                    self.logger.log(f"获取 {date_str} 数据失败: {str(e)}")
                    sse_data.append(0.0)
                    szse_data.append(0.0)

            # 创建DataFrame（确保初始化）
            df = pd.DataFrame({
                'date_str': dates_str,
                'sse_turnover': sse_data,
                'szse_turnover': szse_data
            })

            # 转换日期类型
            if not df.empty:
                df['date'] = pd.to_datetime(df['date_str'], format='%Y%m%d')
                df = df.drop(columns=['date_str'])
                df['total_turnover'] = df['sse_turnover'] + df['szse_turnover']
                print(df['total_turnover'])
                print(df['total_turnover'])
                print(df['szse_turnover'])
                df = df.sort_values('date')
            else:
                df = pd.DataFrame(columns=['date', 'sse_turnover', 'szse_turnover', 'total_turnover'])

            # 过滤日期范围
            mask = (df['date'] >= pd.Timestamp(start_date)) & (df['date'] <= pd.Timestamp(end_date))
            return df.loc[mask]

        except Exception as e:
            import traceback
            self.logger.log(f"成交额获取失败详情:\n{traceback.format_exc()}")
            return pd.DataFrame(columns=['date', 'sse_turnover', 'szse_turnover', 'total_turnover'])



        
    def fetch_sse_turnover(self, date_str: str) -> float:
        """获取上交所成交额（亿元），带重试机制"""
        max_retries = 3
        retry_delay = 1  # 秒

        for attempt in range(max_retries):
            try:
                df = ak.stock_sse_deal_daily(date=date_str)
                if df is None:
                    raise ValueError("API返回空数据")
                
                # 验证数据列
                if "单日情况" not in df.columns or "股票" not in df.columns:
                    raise ValueError("API返回数据结构异常")
                
                # 查找成交金额行（注意繁体字）
                amount_row = df[df["单日情况"] == "成交金额"]
                if amount_row.empty:
                    print(f"上交所 {date} 未找到'成交金额'数据")
                    return 0.0

                value = amount_row["股票"].values[0]
                if not isinstance(value, (int, float)):
                    raise ValueError(f"非数值类型: {type(value)}")
                return round(float(value), 2)
                
            except Exception as e:
                if attempt < max_retries - 1:
                    self.logger.log(f"上交所{date_str} 第{attempt+1}次重试...")
                    time.sleep(retry_delay * (attempt + 1))
                    continue
                self.logger.log(f"获取上交所 {date_str} 成交额失败: {str(e)}")
                return 0.0

    def fetch_szse_turnover(self, date_str: str) -> float:
        """获取深交所成交额（亿元），带数据校验"""
        max_retries = 2
        retry_delay = 2  # 秒

        for attempt in range(max_retries):
            try:
                df = ak.stock_szse_summary(date=date_str)
                if df is None:
                    raise ValueError("API返回空数据")

                # 验证必要列存在
                required_cols = {"证券类别", "成交金额"}
                if not required_cols.issubset(df.columns):
                    missing = required_cols - set(df.columns)
                    raise ValueError(f"缺失列: {missing}")

                # 筛选股票数据
                stock_row = df[df["证券类别"] == "股票"]
                if stock_row.empty:
                    return 0.0

                # 提取并转换数据
                amount_yuan = stock_row["成交金额"].values[0]
                if not isinstance(amount_yuan, (int, float)):
                    raise ValueError(f"非数值类型: {type(amount_yuan)}")

                return round(float(amount_yuan) / 100000000, 2)
                
            except Exception as e:
                if attempt < max_retries - 1:
                    self.logger.log(f"深交所{date_str} 第{attempt+1}次重试，错误: {str(e)}")
                    time.sleep(retry_delay * (attempt + 1))
                    continue
                self.logger.log(f"获取深交所 {date_str} 成交额失败: {str(e)}")
                return 0.0